import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms

import time

# 将图像数据进行归一化处理，使得模型更容易收敛
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True, transform=transform)
test_loader = DataLoader(test_data, batch_size=64, shuffle=True)

print(f"训练集的长度为：{len(train_data)}")
print(f"测试集的长度为：{len(test_data)}")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

drop = 0.4

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.model1 = Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Flatten(),
            nn.Linear(in_features=64 * 4 * 4, out_features=64),
            nn.ReLU(),
            nn.Dropout(drop),
            nn.Linear(in_features=64, out_features=10)
        )


    def forward(self, x):
        x = self.model1(x)
        return x

model = Model()
print(model)
if torch.backends.mps.is_available():
    model = model.to(device)

loss_fn = nn.CrossEntropyLoss()
if torch.backends.mps.is_available():
    loss_fn = loss_fn.to(device)
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
epochs = 40
train_step = 0

# tensorboard记录一下训练时的 loss 变换
writer = SummaryWriter("./logs-train")
for epoch in range(epochs):
    start_time = time.time()
    print(f"-------------第{epoch}轮训练开始了------------")
    model.train()
    train_accuracy = 0
    for data in train_loader:
        images, labels = data
        if torch.backends.mps.is_available():
            images = images.to(device)
            labels = labels.to(device)
        out_labels = model(images)
        loss = loss_fn(out_labels, labels)
        train_accuracy += (out_labels.argmax(dim=1) == labels).sum()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_step += 1
        if train_step % 100 == 0:
            end_time = time.time()
            print(f'训练 100 次的时间为：{end_time - start_time}')
            print(f'-------------训练次数：{train_step}，loss 为: {loss.item()}------------')
            writer.add_scalar('train_loss', loss.item(), train_step)

    writer.add_scalar('train_accuracy', train_accuracy / len(test_data), epoch)
    print(f'epoch：{epoch}, 训练集上的accuracy: {train_accuracy / len(train_data) * 100:.2f}%')

    model.eval()
    test_loss = 0
    test_accuracy = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            if torch.backends.mps.is_available():
                images = images.to('mps')
                labels = labels.to('mps')
            out_labels = model(images)
            test_loss += loss_fn(out_labels, labels)
            test_accuracy += (out_labels.argmax(dim=1) == labels).sum()

    writer.add_scalar('test_loss', test_loss / len(test_data), epoch)
    writer.add_scalar('test_accuracy', test_accuracy / len(test_data), epoch)
    print(f"------------测试集上的 loss: {test_loss / len(test_data)}, accuracy: {test_accuracy / len(test_data)}")



